// Copyright (c) 2010, Google Inc.
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//     * Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//     * Redistributions in binary form must reproduce the above
// copyright notice, this list of conditions and the following disclaimer
// in the documentation and/or other materials provided with the
// distribution.
//     * Neither the name of Google Inc. nor the names of its
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

// ---
// Author: Jeffrey Junfeng Pan (panjunfeng@gmail.com)

// class AssociationRules is used for generating association rules from
// Frequent Itemsets and storing the result rules.

#ifndef FP_GROWTH_ASSOCIATION_RULES_H_
#define FP_GROWTH_ASSOCIATION_RULES_H_

#include <string>
#include <vector>

#include "base.h"
#include "common/scoped_ptr.h"
#include "common/strutil.h"
#include "itemsets.h"

namespace frequent_pattern_mining {

// ------------------------------------------------------------
//                      Declaration
// ------------------------------------------------------------

template <typename T>
class AssociationRules {
 public:
  typedef vector<AssociationRule<T> > AssociationRuleVector;

  // Constructor.
  AssociationRules();

  // Destructor.
  ~AssociationRules();

  // Build the association rules from frequent itemsets.
  void BuildRules(const Itemsets<T>& itemsets,
                  int min_support_count,
                  double min_support_ratio,
                  int min_confidence_count,
                  double min_confidence_ratio);

  // Retrieve the actual association rules.
  const AssociationRuleVector& rules() const;

  // Convert the whole association rules to a string.
  string ToDebugString();

 private:

  // Basic steps for building the association rules.
  void BuildRulesBasicStep(const Itemsets<T>& itemsets,
                           int min_support_count,
                           double min_support_ratio,
                           int min_confidence_count,
                           double min_confidence_ratio,
                           vector<AssociationRule<T> >* rules);

  // Recursive steps for building association rules.
  void BuildRulesRecursiveStep(const Itemsets<T>& itemsets,
                               const ItemsetFrequencyPairBase<T>& lhs,
                               const vector<int>& parents,
                               int min_confidence_count,
                               double min_confidence_ratio,
                               vector<AssociationRule<T> >* rules);

  // The actual association rules.
  scoped_ptr<AssociationRuleVector> rules_;

  // The number of transactions in the orignal dataset.
  int num_transactions_;

  // Min support count.
  int min_support_count_;

  // Min support ratio.
  double min_support_ratio_;

  // Min confidence count.
  int min_confidence_count_;

  // Min confidence ratio.
  double min_confidence_ratio_;

  DISALLOW_EVIL_CONSTRUCTORS(AssociationRules);
};

// ------------------------------------------------------------
//                      Definition
// ------------------------------------------------------------

// Constructor.
template <typename T>
AssociationRules<T>::AssociationRules()
  : rules_(new vector<AssociationRule<T> >) {
}

// Destructor.
template <typename T>
AssociationRules<T>::~AssociationRules() {
}

// Retrieve the actual association rules.
template <typename T>
const typename AssociationRules<T>::AssociationRuleVector&
  AssociationRules<T>::rules() const {
  return *rules_;
}
// Build association rules from frequent itemsets.
template <typename T>
void AssociationRules<T>::BuildRules(
    const Itemsets<T>& itemsets,
    int min_support_count,
    double min_support_ratio,
    int min_confidence_count,
    double min_confidence_ratio) {
  BuildRulesBasicStep(itemsets,
                      min_support_count,
                      min_support_ratio,
                      min_confidence_count,
                      min_confidence_ratio,
                      rules_.get());
}

// Internal function for building association rules: basic step.
template <typename T>
void AssociationRules<T>::BuildRulesBasicStep(
    const Itemsets<T>& itemsets,
    int min_support_count,
    double min_support_ratio,
    int min_confidence_count,
    double min_confidence_ratio,
    vector<AssociationRule<T> >* rules) {
  // Check which argument has a valid value.
  const bool min_support_count_supplied = (min_support_count >= 0);
  const bool min_support_ratio_supplied = (min_support_ratio >= 0);
  const bool min_confidence_count_supplied = (min_confidence_count >= 0);
  const bool min_confidence_ratio_supplied = (min_confidence_ratio >= 0);

  // Only one available valid min_support, either count or ratio.
  CHECK(min_support_count_supplied != min_support_ratio_supplied)
    << ": One and only one of min_support_count "
    << "and min_support_ratio must be set";

  // Only one available valid min_confidence, either count or ratio.
  CHECK(min_confidence_count_supplied != min_confidence_ratio_supplied)
    << ": One and only one of min_confidence_count "
    << "and min_confidence_ratio must be set";

  // Set the number of transactions
  num_transactions_ = itemsets.num_transactions();

  // Empty the existing Association Rules if any.
  rules->clear();

  // Set min_support_count
  min_support_count = min_support_count >= 0 ? min_support_count :
    GetCountFromRatio(num_transactions_, min_support_ratio);

  // Build Association Rules.
  for (typename Itemsets<T>::ItemsetFrequencyPairVector::const_iterator
       i = itemsets.itemsets().begin();
       i != itemsets.itemsets().end();
       ++i) {
    // Construct the trival association rule: lhs ==> (null)
    const ItemsetFrequencyPair<T>& value = *i;
    int threshold = min_confidence_count >= 0 ? min_confidence_count :
      GetCountFromRatio(value.frequency, min_confidence_ratio);
    if (value.frequency >= min_support_count && value.frequency >= threshold) {
      ItemsetFrequencyPairBase<T> lhs;
      lhs.itemset = value.itemset;
      lhs.frequency = value.frequency;
      AssociationRule<T> rule;
      rule.lhs = lhs;
      rule.rhs = lhs;
      rules->push_back(rule);
      // Recursively construct the non-trival association rules: lhs ==> rhs
      BuildRulesRecursiveStep(itemsets,
                              lhs,
                              i->parents,
                              min_confidence_count,
                              min_confidence_ratio,
                              rules);
    }  // if value.frequency >= threshold
  }  // for i (itemsets)

  // Remove redundant items in rhs that appear in lhs.
  for (typename AssociationRuleVector::iterator i = rules->begin();
       i != rules->end();
       ++i) {
  // for (int i = 0; i < rules->size(); ++i) {
    vector<T>& lhs = i->lhs.itemset;
    vector<T>& rhs = i->rhs.itemset;
    for (typename vector<T>::iterator j = rhs.begin();
         j != rhs.end();
         ++j) {
      if (find(lhs.begin(), lhs.end(), *j) != lhs.end()) {
        rhs.erase(j);
        --j;
      }  // if
    }  // for j
  }  // for i (rules)

  // Search and set rhs itemset frequency and combined frequency.
  for (typename AssociationRuleVector::iterator i = rules->begin();
       i != rules->end();
       ++i) {
    if (i->rhs.itemset.size() == 0) {
      i->combined_frequency = i->lhs.frequency;
      i->rhs.frequency = num_transactions_;
    } else {
      bool found = false;
      typename Itemsets<T>::ItemsetFrequencyPairVector::const_iterator index;
      for (typename Itemsets<T>::ItemsetFrequencyPairVector::const_iterator
           j = itemsets.itemsets().begin();
           j != itemsets.itemsets().end();
           ++j) {
        if (itemsets.Equals(j->itemset, i->rhs.itemset)) {
          found = true;
          index = j;
          break;
        }  // if equals
      }  // for j (itemsets)
      CHECK_EQ(true, found) << "rhs itemset must exist. something is wrong.";
      i->combined_frequency = i->rhs.frequency;
      i->rhs.frequency = index->frequency;
    }  // if rhs is not empty
  }  // for i (rules)
}

// Internal function for building association rules: recursive step.
template <typename T>
void AssociationRules<T>::BuildRulesRecursiveStep(
    const Itemsets<T>& itemsets,
    const ItemsetFrequencyPairBase<T>& lhs,
    const vector<int>& parents,
    int min_confidence_count,
    double min_confidence_ratio,
    vector<AssociationRule<T> >* rules) {
  // TODO(panjf) recursive function may be hard to debug and track. avoid using
  // recursive call if possible.
  for (vector<int>::const_iterator i = parents.begin();
       i != parents.end();
       ++i) {
    const ItemsetFrequencyPair<T>& value = itemsets.itemsets()[*i];
    int threshold = min_confidence_count >= 0 ? min_confidence_count :
      GetCountFromRatio(lhs.frequency, min_confidence_ratio);
    if (value.frequency >= threshold) {
      AssociationRule<T> rule;
      rule.lhs = lhs;
      rule.rhs.itemset = value.itemset;
      rule.rhs.frequency = value.frequency;
      rules->push_back(rule);
      BuildRulesRecursiveStep(itemsets,
                              lhs,
                              itemsets.itemsets()[*i].parents,
                              min_confidence_count,
                              min_confidence_ratio,
                              rules);
    }  // if value.frequency >= threshold
  }  // for i
}

// Convert the association rules into a string.
template <typename T>
string AssociationRules<T>::ToDebugString() {
  string result;
  for (typename AssociationRuleVector::const_iterator i = rules_->begin();
       i != rules_->end();
       ++i) {
    // Construct the itemset string of lhs
    string lhs_str;
    const vector<T>& lhs = i->lhs.itemset;
    for (typename vector<T>::const_iterator j = lhs.begin();
         j != lhs.end();
         ++j) {
      lhs_str += ItemToString(*j);
      lhs_str += (j + 1 != lhs.end()) ? " " : "";
    }
    // Construct the itemset string of rhs
    string rhs_str;
    const vector<T>& rhs = i->rhs.itemset;
    for (typename vector<T>::const_iterator j = rhs.begin();
         j != rhs.end();
         ++j) {
      rhs_str += ItemToString(*j);
      rhs_str += (j + 1 != rhs.end()) ? " " : "";
    }
    // Calculate support, confidence and point-wise mutual information
    float support = static_cast<float>(i->lhs.frequency) /
                    static_cast<float>(num_transactions_) * 100.0;
    float confidence = static_cast<float>(i->combined_frequency) /
                       static_cast<float>(i->lhs.frequency) * 100.0;
    float pmi = static_cast<float>(i->combined_frequency) /
                (static_cast<float>(i->lhs.frequency) +
                 static_cast<float>(i->rhs.frequency) -
                 static_cast<float>(i->combined_frequency));
    result += StringPrintf("[%5d / %5d / %5d], %5.3f, "
                           "(%5.1f%% / %5.1f%%), RULE %s ==> %s\n",
                           i->lhs.frequency,
                           i->combined_frequency,
                           i->rhs.frequency,
                           pmi,
                           support,
                           confidence,
                           lhs_str.c_str(),
                           rhs_str.c_str());
  }
  return result;
}
}  // namespace frequent_pattern_mining

#endif  // FP_GROWTH_ASSOCIATION_RULES_H_
